"""
Description: A simple teleoperation example using keyboard input to control a robot in simulation.

Author: Jin Liu
Date: 2025/06/18
"""

from ezgl import tb, Matrix4x4
from ezrobot.sim import RobotClient
from ezrobot.utils import GraphRunner
from ezrobot.motion import OnlineMotion
from ezrobot.motion.mpc_controller.mpc_motion import MPCMotion
from ezrobot.motion.mpc_controller.model import make_fk_function
import random

fk_func = make_fk_function()

if __name__ == '__main__':
    client = RobotClient()
    client.clear_items()
    client.toggle_traj(True)
    target_pose = client.get_pose()
    client.add_pose_axis(target_pose * Matrix4x4.fromScale(4, 4, 4))
    motion = MPCMotion(client)
    
    
    def show_est_point():
        if motion.solver.est_all_q is not None:
            n_dof = motion.solver.nq
            for i in range(5):
                q_i = motion.solver.est_all_q[i*n_dof: (i+1)*n_dof]
                pos, R = fk_func(q_i)
                pose = Matrix4x4.fromRotTrans(R.full(), pos.full().flatten())
                client.set_named_axis(i, pose)

    # 异步运行运动跟踪
    motion_runner = (
        GraphRunner()
        .add_node(0, motion.get_joint)
        .add_node(1, client.set_joint, -1)
        .add_node(2, show_est_point)
        .run_async(frequency=60)
    )
    
    def move_rand():
        x = random.uniform(-0.04, 0.04)   # 范围：-2 到 2
        y = random.uniform(-0.04, 0.04)   # 范围：-2 到 2
        z = random.uniform(0, 0.04)    # 范围：0 到 2
        target_pose.translate(x, y, z, local=False)
        motion.add_pose(target_pose)
        client.set_pose_axis(target_pose * Matrix4x4.fromScale(4, 4, 4))

    def move(x, y, z):
        target_pose.translate(x, y, z, local=False)
        motion.add_pose(target_pose)
        client.set_pose_axis(target_pose * Matrix4x4.fromScale(4, 4, 4))
        print(target_pose.xyz)
    
    with tb.window("control", None, 10, size=(500, 500), pos=(400, 0), frameless=False):
        tb.add_text("text", text="Sim Control")
        tb.add_stretch(1)
        tb.add_key_binding(tb.Key.Key_D, lambda: move(0.01, 0, 0))
        tb.add_key_binding(tb.Key.Key_A, lambda: move(-0.01, 0, 0))
        tb.add_key_binding(tb.Key.Key_W, lambda: move(0, 0.01, 0))
        tb.add_key_binding(tb.Key.Key_S, lambda: move(0, -0.01, 0))
        tb.add_key_binding(tb.Key.Key_Q, lambda: move(0, 0, 0.01))
        tb.add_key_binding(tb.Key.Key_E, lambda: move(0, 0, -0.01))
        tb.add_key_binding(tb.Key.Key_Space, lambda: move_rand())

    tb.exec()